
import numpy as np
from Alexnet import alexnet2
from random import shuffle
import pandas as pd

# what to start at
START_NUMBER = 60

# what to end at
hm_data = 100

# use a previous model to begin?
START_FRESH = False

WIDTH = 160
HEIGHT = 86
LR = 5e-3
EPOCHS = 5
MODEL_NAME = 'AlexModel/Alexmodel.model'
EXISTING_MODEL_NAME = 'AlexModel/Alexmodel.model'
file_name = 'train_data.npy'

model = alexnet2(WIDTH, HEIGHT, LR)

if not START_FRESH:
    model.load(EXISTING_MODEL_NAME)

for i in range(EPOCHS):
    '''
    data_order = [i for i in range(START_NUMBER,hm_data+1)]
    shuffle(data_order)
    for i in data_order:
        train_data = np.load(file_name,allow_pickle=True)

        df = pd.DataFrame(train_data)
        df = df.iloc[np.random.permutation(len(df))]
        train_data = df.values.tolist()
    '''
    train_data = np.load(file_name, allow_pickle=True)
    train = train_data[:-4500]
    test = train_data[-4500:]

    X = np.array([i[0] for i in train]).reshape(-1, WIDTH, HEIGHT, 1)
    Y = [i[1] for i in train]

    test_x = np.array([i[0] for i in test]).reshape(-1, WIDTH, HEIGHT, 1)
    test_y = [i[1] for i in test]

    model.fit({'input': X}, {'targets': Y}, n_epoch=1, validation_set=({'input': test_x}, {'targets': test_y}),
              snapshot_step=2500, show_metric=True, run_id=MODEL_NAME, shuffle=True)

    model.save(MODEL_NAME)